'''
Basic class to submit jobs on Philly through
the REST API.
'''

import uuid
import json
import logging
import threading
import time
from urllib.parse import urlencode
from http.client import responses
import requests
import requests_ntlm
from src.vault import Credentials
from src.exceptions import (
    ManagerException, AuthorizationException, ConnectionException, RequestFailed)
import click

LOG = logging.getLogger('philly')
THREAD_LOCAL = threading.local()


def get_ssl_warning():
  """Serializes a warning when ssl is unsuccessful.

  Returns:
    A message string.
  """
  msg = "Your system did snot accept SSL certificate presented by " \
      "https://philly. To fix this, do one of the following:\n" \
      "- configure your system to trust philly's certificate authority.\n" \
      "- from certmgr in windows, export the philly root cert and provide " \
      "the file using --cert or via PHILLY_CERT\n" \
      "- use --insecure, but keep in mind that your password is sent in " \
      "plain text to an untrusted server."
  return msg


class PhillyManager(object):
  """Provides an interface to send REST requests to the Philly server.
  """
  def __init__(self, cluster, vc, custom_docker_name=None,
               docker_registry="phillyregistry.azurecr.io", docker_repo=None,
               docker_tag=None, insecure=False):
    """Constructor of PhillyManager.

    For version 2.0 of the docker images, provide a `docker_repo`, e.g.
    philly/jobs/custom/openai-gym-ale and a `docker_tag`, e.g. tf13-py36.
    Note that using the "latest" tag does not work. To keep using v1,
    provide a `custom_docker_name` as usual, e.g. custom-malibu-tf-13-py36.

    Args:
      cluster: physical cluster to use.
      vc: virtual cluster the user belongs to.
      custom_docker_name: docker name for philly dockers version 1 (v1).
      docker_registry: registry for dockers version 2 (v2).
      docker_repo: repository for dockers version 2 (v2).
      docker_tag: tag for dockers version 2 (v2).
      insecure: proceed for insecure authentication.
    """
    LOG.debug("Philly Manager, cluster=%s, vc=%s", cluster, vc)
    cred = Credentials.get()
    self.username = cred.username
    self.password = cred.password
    if cred.cert is not None:
      self.ssl_verify_peer = cred.cert
    else:
      self.ssl_verify_peer = True

    if cred.insecure:
      self.ssl_verify_peer = False

    self.custom_docker_name = custom_docker_name
    self.docker_registry = docker_registry
    self.docker_repo = docker_repo
    self.docker_tag = docker_tag
    self.vc = vc
    # physical cluster
    self.cluster = cluster

    self.domain_joined = False
    if cred.domain_joined:
      self.domain_joined = True
      self.curl = self.ms_curl

  @property
  def session(self):
    # initialize auth once (before spawning threads)
    session = getattr(THREAD_LOCAL, 'session', None)
    if session is not None:
      age = time.time() - THREAD_LOCAL.session_time
      if age > 180:
        # invalidate session after 3 minutes to prevent timeouts
        session = None
    if session is None:
      THREAD_LOCAL.session = requests.Session()
      THREAD_LOCAL.session.auth = requests_ntlm.HttpNtlmAuth(self.username, self.password)
      THREAD_LOCAL.session.get('https://philly', verify=self.ssl_verify_peer)
      THREAD_LOCAL.session_time = time.time()
      session = THREAD_LOCAL.session
    return session

  def _co_initialize(self):
    # to use native ntlm auth on windows, the CoInitialize function *has to* be
    # called *in every thread*.
    #
    # What does it do? It "initializes the COM libraries for the calling thread".
    co_initialized = getattr(THREAD_LOCAL, 'co_initialized', False)
    if not co_initialized:
      import pythoncom
      pythoncom.CoInitialize()
      THREAD_LOCAL.co_initialized = True

  def abort_job(self, job_id):
    """Sends an abort request.

    Args:
      job_id: the appID of the job to be aborted.

    Returns:
      The json result from the server.
    """
    res = self.curl('abort', response_fmt='json',
                    clusterId=self.cluster, jobId=job_id)
    LOG.debug(res)
    return res

  def abort_all_jobs(self):
    """Aborts all jobs pertaining to the current user.
    """
    active_jobs = self.get_active_jobs()
    if active_jobs is None:
      return
    LOG.debug("Found %d active jobs: %s", len(active_jobs), active_jobs)
    if len(active_jobs) > 0:
      for active_job in active_jobs:
        self.abort_job(active_job)

  def get_job_status(self, job_id):
    """Get status of a job.

    Status can be Running/Failed/Queued/Pass.

    Args:
      job_id: the appID of the job to get the status for.

    Returns:
      The json result from the server containing a 'status' field.
    """
    return self.curl('status', response_fmt='json',
                     clusterId=self.cluster, jobType='cust',
                     content='partial', vcId=self.vc, jobId=job_id)

  def ms_curl(self, func, timeout=None, response_fmt='text', **kwargs):
    self._co_initialize()
    func = func + "?" + urlencode(kwargs)
    LOG.debug('PHILLY API Call: %s', func)
    url = "https://philly/api/%s" % func
    import win32com.client
    COM_OBJ = win32com.client.Dispatch('WinHTTP.WinHTTPRequest.5.1')
    COM_OBJ.SetAutoLogonPolicy(0)
    COM_OBJ.Open('GET', url, False)
    COM_OBJ.Send()
    if response_fmt == 'binary':
      return COM_OBJ.ResponseBody
    resp = COM_OBJ.ResponseText
    if response_fmt == 'json':
      try:
        return json.loads(resp)
      except json.JSONDecodeError as e:
        raise RequestFailed(func, resp, str(e))
    if response_fmt == 'text':
      return resp
    raise ValueError("Unknown response_fmt: '%s'" % response_fmt)

  def curl(self, func, response_fmt='text', timeout=None, **kwargs):
    """Helpers for submitting REST requests.

    Returns:
      The json result from the server.
    """
    func = func + "?" + urlencode(kwargs)
    LOG.debug('PHILLY API Call: %s', func)
    try:
      resp = self.session.get("https://philly/api/%s" % func, timeout=timeout,
                              verify=self.ssl_verify_peer)
    except requests.exceptions.SSLError:
      raise ManagerException(get_ssl_warning())
    except requests.exceptions.ConnectionError as e:
      raise ConnectionException(
          "Connection error. There seems to be a problem reaching "
          "https://philly. Make sure you are connected to CorpNet. "
          "Detailed message: %s" % (e,))
    if resp.status_code == 401:
      raise AuthorizationException("Authorization error.")
    if resp.status_code != 200:
      raise RequestFailed(func, "Status %d: %s" % (
          resp.status_code, responses[resp.status_code]), "")
    if response_fmt == 'binary':
      return resp.content
    if response_fmt == 'text':
      return resp.content.decode('utf-8')
    if response_fmt != 'json':
      raise ValueError("Unknown response_fmt: '%s'" % response_fmt)
    try:
      return resp.json()
    except json.JSONDecodeError as e:
      raise RequestFailed(func, resp.content, str(e))

  def get_logs(self, job_id, log_type="stdout"):
    """Get stdout for `job_id`.

    Args:
      job_id: the id of the job.
      log_type: which type of log to get.

    Returns:
      A string containing the stdout/stderr of the job.
    """
    return self.curl('log', vcId=self.vc, jobId=job_id,
                     clusterId=self.cluster, jobType='cust',
                     logType=log_type, logRev='latest',
                     content='full')

  def get_active_jobs(self):
    """Gets a list of active jobs for the user.

    Returns:
      A list of appID corresponding to active jobs.
    """
    res = self.curl('list', response_fmt='json', clusterId=self.cluster, vcId=self.vc)
    jobs = res['runningJobs'] + res['queuedJobs']

    active_jobs = []
    for job in jobs:
      if job['username'] == self.username:
        active_jobs.append(job['appID'])
    return active_jobs

  def submit_job(self, job_path, job_name, data_dir, n_gpu=1, args=None, dry_run=False, submit_args=None):
    """Submits a job.

    Runs the executable specified in `job_path`, with name `job_name`
    using extra arguments as specified in `args` and `data_dir`.

    Args:
      job_path: script path relative to the user home in the cluster.
      job_name: description for this run.
      data_dir: path to the data in the cluster.
      n_gpu: number of gpus to use for this job.
      args: argument string that need to be passed to this script.
      submit_args: custom args to submit to PhillyAPI

    Returns:
      The json result from the server containing appID of the submitted job.

    Raises:
      ManagerException: for authorization errors or unsuccessful submission.
    """
    assert type(n_gpu) == int and n_gpu > 0
    assert len(job_name) != 0, 'Must specify a job description!'
    LOG.debug('Submitting %s: %s, n_gpu = %s, args = %s, docker version = %s',
              job_name, job_path, n_gpu, args,
              "1" if self.custom_docker_name else "2")
    if not data_dir:
      raise ValueError("Expected a valid remote data location.")
    if self.custom_docker_name:
      postfields = dict(customDockerName=self.custom_docker_name)
    else:
      postfields = dict(Registry=self.docker_registry,
                        Repository=self.docker_repo,
                        Tag=self.docker_tag)
      postfields.update(
          dict(buildId='0000',
               toolType='cust',
               clusterId=self.cluster,
               vcId=self.vc,
               configFile=job_path,
               minGPUs=n_gpu,
               name="pt-{}".format(job_name),
               IsDebug='false',
               IsMemCheck='false',
               OneProcessPerContainer='true',
               numFinishedJobs=1,
               IsCrossRack='false',
               inputDir=data_dir,
               userName=self.username))
    if args is not None:
      postfields["extraParams"] = args
    if submit_args is not None:
      postfields.update(submit_args)
    if not dry_run:
      res = self.curl("submit", response_fmt='json', **postfields)
    else:
      res = dict(jobId=str(uuid.uuid4().hex) + "_DRY_RUN")
    if 'error' in res:
      raise ManagerException("an error occured while submitting "
                             "the job: %s" % (res['error'],))
    else:
      LOG.debug("job submission success %s", res['jobId'])
    return res

  @staticmethod
  def create_from_config(cfg):
    """Factory method for creating a `PhillyManager`.

    Args:
      cfg: a `PTConfig` object.

    Returns:
      A `PhillyManager` object.
    """
    philly_mg = PhillyManager(docker_repo=cfg.auth.docker_repo,
                              docker_tag=cfg.auth.docker_tag,
                              vc=cfg.auth.vc, cluster=cfg.auth.cluster)
    return philly_mg

if __name__ == '__main__':
  import os
  user = 'v-zhhe'
  cert = 'data\\ms-internal-root.crt'
  Credentials.instantiate(username=user, insecure=False, cert=cert,
                          domain_joined=True)
  cred = Credentials.get(True)
  print(cred.password)
  job_mgr = PhillyManager(vc='msrlabs', cluster='gcr')
  result_json = job_mgr.curl('summary', response_fmt='json', clusterId='gcr')
  print(result_json)